import baostock as bs
from tqdm import tqdm
import pandas as pd
import datetime

# max concurrent process number
CONCURRENCY = 10
# information columns needed
FIELDS = ("date,code,open,high,low,close," +
          "preclose,volume,amount,adjustflag,turn," +
          "tradestatus,pctChg,peTTM,pbMRQ,psTTM,pcfNcfTTM,isST")

def get_stock_list():
    #return stock list of the component of csi500 index
    rs_500 = bs.query_zz500_stocks()
    rs_50 = bs.query_sz50_stocks()
    rs_300 = bs.query_hs300_stocks()
    stock_list = []
    while (rs_500.error_code == '0') & rs_500.next():
        stock_list.append(rs_500.get_row_data())
    while (rs_50.error_code == '0') & rs_50.next():
        stock_list.append(rs_50.get_row_data())
    while (rs_300.error_code == '0') & rs_300.next():
        stock_list.append(rs_300.get_row_data())
    return stock_list

def get_one_stock(code: str):
    # get single stock data
    rs = bs.query_history_k_data_plus(
        code=code,
        fields=FIELDS,
        start_date='2025-03-01', end_date='2025-06-01',
        frequency="d", adjustflag="3")
    # transformer into dataframe
    data_list = []
    while (rs.error_code == '0') & rs.next():
        data_list.append(rs.get_row_data())
    result = pd.DataFrame(data_list, columns=rs.fields)
    return result

def main():
    lg = bs.login()

    stock_list = get_stock_list()

    list_df = []
    for item in tqdm(stock_list[0:10]):
        list_df.append(get_one_stock(code=item[1]))
    df = pd.concat(list_df, axis=0)

    numerical_columns = list(df.columns)
    numerical_columns.remove("date")
    numerical_columns.remove("code")
    numerical_columns.remove("isST")
    numerical_columns.remove("tradestatus")

    for item in numerical_columns:
        df[item] = pd.to_numeric(df[item])

    for item in ["isST", "tradestatus"]:
        df[item] = df[item].astype(int).astype(bool)

    df['date'] = pd.to_datetime(df['date'])
    df['insert_time'] = pd.to_datetime(datetime.datetime.now())
if __name__ == "__main__":
    main()